#!/usr/bin/env python
'''todo
report more after paramtransfer
param average by structure match
set resname by atom match
combine frag with multiple dihedrals
change pair_nb when fit torsion
list res lib and pickle them for later use
paramrefine do not need log file contain internal coords
internal coords should be generated in coord magic
build whole system
support libpargen and intermol
preview a torsion profile
partial param transfer
eq param transfer
'''

from ztop import *
import os
import sys
import argparse
import coordmagic as cm
import shutil


parser=argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
                               description='refine topology parameters by gaussian fchk and log file\n'
                               'combine small molecules to large ones')
parser.add_argument("-g",dest='gentop',metavar='gen_top',default='',type=str,
                    help='generate top from a structure file, by third party program \n'
                         'the format and some default value is:\n'
                         '-g "input.log;charge=0;spin=1;ff=gaff2;q=resp;ffp=amber;qp=multiwfn;charge_only"\n'
                         'default forcefield is gaff2 and default q is resp by multiwfn\n'
                         'recommanded file formats are: log, mol2, pdb\n'
                         'supported ff(force field) are:\n'
                         'gaff2, gaff, amber(requires antechamber)\n'
                         'opls-aa (requires libpargen, not available for now)  '
                         'SMIRNOFF (requires openff, not available for now) \n'
                         'available q(charge methods) are: \n'
                         'resp resp2 am1-bcc gasteiger cm5 ddec3 ddec6 \n'
                         'ffp and qp are program for parameter generation and change generateion\n'
                         'Normally they will set automatically according to charge type and ff type\n')
parser.add_argument("-f",dest='fragment',metavar='fragment',action='append',default=[],
                    help='Define a fragment, the formats are:'
                         'D;p=ABC.top;x=ABC.gro;site=D1:2-3,D2:11-23;resname=DON;comment="any words"\n'
                         'The first letter D is the label of this residue used only in this program\n'
                         'p=,x= set the top and coord file of the residue\n'
                         'site= define cutting site of the residue. D1 and D2 are site label,\n'
                         'The second to last character of site label are used for site match\n'
                         'following the shortest match rule: e.g. D12 will match A1 or A12 but not A13\n'
                         'X542 will match Y5,Y54,Y542,Y5421 but not X6, X55, X543'
                         'D1:2-3 define site by breaking the bond 2-3 where atom 2 belongs to the residue.')
parser.add_argument("-b",dest='build',metavar='build',default='',type=str,
                    help = 'set a list of res label seperated by - to build the molecule. e.g.\n'
                    'H[AD]5T is a six mer\n'
                    'C[L]3[M]6[N]12 is a dendrimer of three generation\n'
                    'Note that you should first define residue by -r option')
parser.add_argument("-p",dest='top',metavar='initial_top',default='',type=str,
                    help='initital top file\n'
                         '"-p ABC.prmtop" or "-p ABC.top"')
parser.add_argument("-x",dest='coord',metavar='initial_coord',default='',type=str,
                    help='initial coord file\n'
                         'support format:\n'
                         'pbd, mol2, gro, inpcrd, gro')
parser.add_argument("-r",dest='refine',metavar='refine_param',default='',type=str,
                    help='refine parameters, the format is:\n'
                          '-r "e=ABC.log;k=ABC.fchk;c=ABC.chg" to refine equilibrium bond and angle by ABC.log\n'
                          'and bond and angle force constant by ABC.fchk with Modified seminario method\n'
                          'and charge from charge file ABC.chg\n'
                          '-r "e;k;c" will refine params from default file\n')
parser.add_argument("-d",dest='dihedral',metavar='dihedral_fit',default='',type=str,
                    help='fit a torsion parameter, the format is:\n'
                         '-d "15-16;angle=e5,180,0;period=1,2,3,4;eth=60;opt=no;\n'
                         'ar=-30~30,150~180;nd=3;kth=0.02;regen=gjf"\n'
                         '15-16 specify a bond to rotate and dihedral of which to fit\n'
                         'angle:e5,180,0 rotate dihedral from 0 to 180 with increment of 5\n'
                         'period:1,2,3,4 fit the dihedral with period of 1 2 3 and 4, default is auto\n'
                         'eth:60 fit the point with energy lower than 60kJ/mol\n'
                         'opt:no do not add constrain information to the end of generated gjf file\n'
                         'opt:fixone fix one dihedral, and opt:fixall fix all dihedrals\n'
                         'ar:-30~30 select angle range from -30 to 30 to fit\n'
                         'The purpose of angle range is to aviod large steric to intefere torsion param\n'
                         'nd:1,2,3 indicate only fit parameters for the first three dihedrals\n'
                         'remember add opt(modredundant) to the keywords of gjf by --gjf options\n'
                         'kth:0.02 indicate that k value below this value (unit: kCal/mol)\n'
                         'regen:gjf will regenerate gjf file and regen:xyz will regenerate xyz file from log\n'
                         'The program will first check a xyz file named ABC_DFit_15-16 in current directory\n'
                         'where the comment line is the energy of the structure in kJ/mol\n'
                         'If the xyz file is not found,\n'
                         'The program will read log files in ABC_TSFit_15-16 and genererate the xyz file\n'
                         'If the directory is not found, the program will create it and gjf files\n'
                         'The user should get the log file by themselves' )
parser.add_argument("-t",dest='transfer',metavar='transfer_param',action='append',default=[],type=str,
                    help = 'transfer parameters from a residue that define by -r\n'
                    'the format is:\n'
                    'd;match=12-45,73,9;a=40,43,9-12;p=d\n'
                    'The D is the residue label\n'
                    'm= designate the atom serial in D that used to match substruct in main structure\n'
                    'a= designate the parameter of which atoms will be transfered.\n'
                    'p= cbadv specifed which paramters will be transfers\n'
                    'c:charge,v:vdw,b:bond,a:angle,d:dihedral')
# parser.add_argument("-a",dest='average',metavar='average_param',default='',type=str,
#                     help = 'average parameters for substructrures\n'
#                        'the format is:\n'
#                        'm=12-45,73,9;a=40,43,9-12;p=d\n'
#                        'm= designate the atom serial used to match substructure\n'
#                        'a= designate the parameter of which parameters will be averaged.\n'
#                        'p= cbadv specifed which paramters will be transfers\n'
#                        'c:charge,v:vdw,b:bond,a:angle,d:dihedral')
parser.add_argument("-o",dest='output',metavar='output',default='',type=str,
                    help='set output filename and file extension. For example:\n '
                         '"-o ABC.top,ABC.gro" generate file for gromacs\n'
                         '"-o ABC.prmtop,ABC.inpcrd" generate file for amber\n')
# parser.add_argument("--resname",dest='resname',metavar='resname',default='',type=str,
#                     help='set the residue name of the structure\n'
#                          'seprate multiple names by comma')
parser.add_argument("--savelib",dest='savelib',action="store_true",
                    help='save the current fragment lib to directory FRAG_LIB\n')
parser.add_argument("--loadlib",dest='loadlib',action="store_true",
                    help='load the fragment lib from directory FRAG_LIB\n')
parser.add_argument("--gjf",dest='gjf',metavar='gjf_option',default='',type=str,
                    help='can set some gjf file options. For example:\n '
                         '"--gjf "nproc:8;mem:4GB;charge:0;spin:1;extra:eps=3.0;vdw:em(gd3bj)\n'
                         'method:b3lyp;basis:def2svp;solvent:none;addkey:addition keywords"\n')
parser.add_argument("--opt",dest='opt',action='store_true',
                    help='If this option is set, then the output geom will optimized by openmm\n')
parser.add_argument("--checkenv",dest='checkenv',action='store_true',
                    help='check the program and python package needed by this program\n')
parser.add_argument('--version', action='version', version='%(prog)s 1.0')
args=parser.parse_args()
# st=cm.read_structure('Dfrag.log')
# st.M.gen_molecules()
# D=st.M.gen_frag_by_breakbond(D1='18-36',D2='4-27')
# A0=st.M.gen_frag_by_breakbond(A2='27-4')
# sta = cm.read_structure('Afrag.log')
# sta.M.gen_molecules()
# A=sta.M.gen_frag_by_breakbond(A1='20-22',A2='14-4')
# D0=sta.M.gen_frag_by_breakbond(D2='22-20')
# frag1 = st.M.connect_frag(D,A,joinsite='D-A')
# # print(frag1)
# frag2 = st.M.connect_frag(frag1,D,joinsite='A-D')
# frag3 = st.M.connect_frag(frag2,A,joinsite='D-A')
# frag4 = st.M.connect_frag(frag3,D,joinsite='A-D')

pr = None
rc = None

if args.checkenv:
    print('check parmed...',end='')
    try:
        import parmed
    except:
        print('\nError!!! python package parmed not available')
    else:
        print('parmed version {:s} detected'.format(parmed.__version__))
    print('check networkx...',end='')
    try:
        import networkx
    except:
        print('\nError!!! python package networkx not available')
    else:
        print('networkx version {:s} detected'.format(networkx.__version__))
    print('check Multiwfn...',end='')
    multiwfn = shutil.which('Multiwfn')
    multiwfn_win = shutil.which('Multiwfn.exe')
    if not (multiwfn or multiwfn_win):
        print('\nError!!! {:s} requires Multiwfn, but it is not available'.format(qfftype))
    else:
        print('success!')
    print('check AmberTools...',end='')
    antechamber = shutil.which('antechamber')
    parmchk2 = shutil.which('parmchk2')
    tleap = shutil.which('tleap')
    p2p = {'antechamber': antechamber,
           'parmchk2': parmchk2,
           'tleap': tleap}
    if not (antechamber and parmchk2 and tleap):
        print('\nError! AmberTools program({:s}) not available'
              .format(','.join([i for i, v in p2p.items() if not v])))
    else:
        print('success!')
    print('check openmm...',end='')
    try:
        import simtk.testInstallation as TI
    except:
        print('Error!!! python package simtk not available')
    else:
        TI.main()

if args.loadlib:
    # print(args.residue)
    rc = ResidueComposer()
    rc.load_lib()
    rc.print_lib()

if len(args.fragment) > 0:
    # print(args.residue)
    rc = ResidueComposer()
    for r in args.fragment:
        rc.parse_residue(r)
    rc.print_lib()

if rc:
    if args.savelib:
        rc.save_lib()
    if args.build:
        mol=rc.compose_mol(args.build)
        final_top = mol['pmd_struct']
        pr = ParamRefine(pmd_struct=final_top)
        pr.set_charge()
        # for d in pr.ps.dihedrals:
        #     if d.type.scee > 100:
        #         print(d)

if args.gentop:
    inpfile = args.gentop.split(';')[0]
    kwarg = {v.split('=')[0]:v.split('=')[1] for v in args.gentop.split(';') if len(v.split('='))==2}
    kwarg['gjf_options'] = args.gjf
    pg = ParamGen(inpfile,**kwarg)
    if "charge_only" in args.gentop:
        pg.gen_param(charge_only=True)
    else:
        top,crd,chg = pg.gen_param()
        pr = ParamRefine(top_file = top)
        pr.load_coord(crd)

if args.top:
    pr = ParamRefine(top_file = args.top)

if pr:
    if args.coord:
        pr.load_coord(args.coord)
    if args.refine:
        rf = args.refine.split(';')
        rf_task = {i.split('=')[0]:'' for i in rf}
        rf_task.update({i.split('=')[0]:i.split('=')[1] for i in rf if len(i.split('=')) == 2})
        if 'c' in rf_task:
            if not rf_task['c']:
                rf_task['c'] = pr.basename + '.chg'
            if not os.path.isfile(rf_task['c']):
                print('Error! charge file {:s} not found in current directory'.format(rf_task['c']))
                sys.exit()
            print('Statistics on charge refinement:')
            pr.refine_charge(rf_task['c'])
        if 'e' in rf_task:
            if not rf_task['e']:
                rf_task['e'] = pr.basename + '.log'
            if not os.path.isfile(rf_task['e']):
                print('Error! log file {:s} not found in current directory'.format(rf_task['e']))
                sys.exit()
        if 'k' in rf_task:
            if not rf_task['k']:
                rf_task['k'] = '{:s}.fchk'.format(pr.basename)
            if not os.path.isfile(rf_task['k']):
                print('Error! fchk file {:s} not found in current directory'.format(rf_task['k']))
                sys.exit()
        if 'e' in rf_task and 'k' not in rf_task:
            st=cm.read_structure(rf_task['e'])
            st.params={'coords':st.coord}
            if not pr.ps.positions:
                pr.load_coord(rf_task['e'])
            coord1,ene1=pr.minimize()
            r1=pr.compare_structure(st.coord,coord1)
            pr.refine_by(st)
            coord2,ene2=pr.minimize()
            r2=pr.compare_structure(st.coord,coord2)
            print('Before parameter refine by reset eq bond and angle (openmm opt vs input):')
            print('Energy: {:.3f} kJ/mol'.format(ene1._value))
            print(r1)
            print('After parameter refine by reset eq bond and angle (openmm opt vs input):')
            print('Energy: {:.3f} kJ/mol'.format(ene2._value))
            print(r2)
        if 'k' in rf_task and 'e' in rf_task:
            ms = ModSeminario(rf_task['e'], rf_task['k'])
            if not pr.ps.positions:
                pr.load_coord(rf_task['e'])
            coord1,ene1=pr.minimize()
            report1=pr.compare_structure(ms.coords,coord1)
            pr.refine_by(ms)
            coord2,ene2=pr.minimize()
            report2=pr.compare_structure(ms.coords,coord2)
            print('Before parameter refine by Modified seminario method (openmm opt vs input):')
            print('Energy: {:.3f} kJ/mol'.format(ene1._value))
            print(report1)
            print('After parameter refine by Modified seminario method (openmm opt vs input):')
            print('Energy: {:.3f} kJ/mol'.format(ene2._value))
            print(report2)
    if len(args.transfer) > 0:
        for t in args.transfer:
            res = t.split(';')[0]
            param = {i.split('=')[0]:i.split('=')[1] for i in t.split(';')[1:]}
            coord1, ene1 = pr.minimize()
            pr.transfer_param_from(rc.res_lib[res],**param)
            coord2, ene2 = pr.minimize()
            report = pr.compare_structure(coord1,coord2)
            print('Energy before transfer: {:f} kJ/mol; Energy after transfer {:f} kJ/mol'
                  .format(ene1._value,ene2._value))
            print('Min structure difference after parameter transfer:')
            print(report)
    if args.dihedral:
         bond = args.dihedral.split(';')[0]
         param = {i.split('=')[0]:i.split('=')[1] for i in args.dihedral.split(';')[1:]}
         if args.gjf:
             param.update({'gjf':args.gjf})
         TorsionFit(pr,bond,**param)
    if args.opt:
         coords,ene = pr.minimize()
         pr.ps.positions = coords
         print('the mm energies after minimize is {:.3f} kJ/mol'.format(ene._value))
    if args.output:
        for f in args.output.split(','):
            pr.adj_dihedral(contract= 'all')
            print("Saving file {:s}...".format(f))
            pr.ps.save(f,overwrite=True)






# Dfrag_amb = pmd.load_file('Dfrag.prmtop','Dfrag.inpcrd')
# # system = Dfrag_amb.createSystem(nonbondedMethod=app.NoCutoff,constraints=app.HBonds)
# prd = ParamRefine(Dfrag_amb)
# ms = ModSeminario("Dfrag.log","Dfrag.fchk")
# prd.refine_by(ms)
# rc = ResidueComposer()
# fragD = rc.gen_res(prd.ps,resname='DON',D1='18-36',D2='4-27')
#
# Afrag_amb = pmd.load_file('Afrag.prmtop','Afrag.inpcrd')
# prd = ParamRefine(Afrag_amb)
# ms = ModSeminario("Afrag.log","Afrag.fchk")
# prd.refine_by(ms)
# rc = ResidueComposer()
# fragA = rc.gen_res(prd.ps,resname='ACC',A1='20-22',A2='14-4')
# fragC = rc.connect_res(fragD,fragA,joinsite='D-A')
# fragC['pmd_struct'].save('testxxx.pdb',overwrite=True)
# fragC['pmd_struct'].save('testxxx.top',overwrite=True)
# pmd_fragC = fragC['pmd_struct']
# prc=ParamRefine(pmd_fragC)
# # prc.get_BAD()
#
# # for dtype in pmd_fragC.dihedral_types:
# #     print(dtype)
# system = pmd_fragC.createSystem(nonbondedMethod=app.NoCutoff,constraints=app.HBonds)
# # Create the integrator to do Langevin dynamics
# integrator = mm.LangevinIntegrator(
#                         300*kelvin,       # Temperature of heat bath
#                         1.0/picoseconds,  # Friction coefficient
#                         2.0*femtoseconds, # Time step
# )
#
# sim = app.Simulation(pmd_fragC.topology, system, integrator)
#
# # Set the particle positions
# sim.context.setPositions(pmd_fragC.positions)
# state = sim.context.getState(getEnergy=True)
# print(state.getPotentialEnergy())
# sim.minimizeEnergy()
# state = sim.context.getState(getEnergy=True)
# positions = sim.context.getState(getPositions=True)
# bad = prc.get_BAD(positions.getPositions(True))
# prc.compare_BAD(bad,bad)
# # app.PDBFile.writeFile(sim.topology, positions, open('output.pdb', 'w'))
# print(state.getPotentialEnergy())

# ms.modified_Seminario_method()
# prd.refine_bond_param(ms.bond_s2p)
# prd.refine_angle_param(ms.angle_s2p)
# dfrag = prd.ps[idx]
# new_idx_d = {i:v[1] for i,v in enumerate(sorted(frag1[1].items(),key=lambda x:x[0])) }
# print(new_idx_d)
# ia,oa,ib,ob = [i-1 for i in frag1[3]]
# print(ia,oa,ib,ob)
# # b=get_bond(prd.ps,[ia,oa])
# a=get_angles(prd.ps,[oa,ia])
# d=get_dihedrals(prd.ps,[oa,ia])
# print(prd.ps.angles)
# print(pr.ps.atoms[0].angle_partners)
# print(pr.ps.atoms[0].list)

# Afrag_amb = pmd.load_file('Afrag.prmtop','Afrag.inpcrd')
# pra = ParamRefine(Afrag_amb)
# ms = ModSeminario("Afrag.log","Afrag.fchk")
# ms.modified_Seminario_method()
# pra.refine_bond_param(ms.bond_s2p)
# pra.refine_angle_param(ms.angle_s2p)
# pra.ps.residues[0].name = 'ACC'
# idx=  '@' + ','.join([str(i) for  i in frag1[2].keys()])
# print(idx)

#
# afrag = pra.ps[idx]
#
# print(new_idx_d)
# print(cfrag.atoms[26].idx)
# bond = get_bond(cfrag,[new_ia_idx,new_ib_idx])
# angle = get_angles(cfrag,[new_ia_idx,new_ib_idx])
# dihedrals = get_dihedrals(cfrag,[new_ib_idx,new_ia_idx])
# for angle in a:
#     a1 = angle.atom1
#     a2 = angle.atom2
#     a3 = angle.atom3
#     # get new index
#     print(angle.type)
# print(bond)
# # print(angle.type)
# print(dihedrals)
# print(cfrag.residues)
# print(dir(amb.bonds[0].atom1))
# print(amb.bonds[0].atom1.atom_type)
# print(amb.bonds[0].atom1.idx)
# prf = ParamRefine('test.top')
# print(amb.angles)
# print(amb.angles)
#print(dir(amb))
# if len(args.output) > 0:
#     for o in args.output:
#         amb.save(o,overwrite=True)

# a0=np.array([0,1,0,1])
# a1=np.array([0,-1,0,1])
# a2 = np.array([1,0,0,1])
# m=st.T.translate_along(a1,a2,1)
# print(a0@m)
# m=st.T.rotate_to_align(a1,a2,a0)
# print(m)
# print(a1@m)
# m=st.T.rotate_cw_around(a0,a1,135)
# print(a2@m)

# m0 = st.T.translate_to_align(a0,[0,0,0,1])
# aaa = st.MM.torsion([3.56994303,1.62762092,-0.87365150],[3.21327019,1.12322273,0.00000000],
#                      [1.67327019,1.12324171,0.00000000],[1.31661627,2.13204999,0.00195600])
# m=st.T.rotate_dihedral_to([3.56994303,1.62762092,-0.87365150],[3.21327019,1.12322273,0.00000000],
#                      [1.67327019,1.12324171,0.00000000],[1.31661627,2.13204999,0.00195600],150)
# newA = np.array([3.56994303,1.62762092,-0.87365150,1])@m
# bbb = st.MM.torsion(newA,[3.21327019,1.12322273,0.00000000],
#                     [1.67327019,1.12324171,0.00000000],[1.31661627,2.13204999,0.00195600])
# print(newA)
# print(aaa)
# print(bbb)
# print(m)
# print(a2@m)

# print(np.linalg.norm(a1[:3]))
# print(np.linalg.norm((a1@a)[:3]))
# print(a1@a)

# prf.read_top()
# prf.param_as2p = ms.angle_s2p
# prf.param_bs2p = ms.bond_s2p
# prf.map_serial()
# prf.compare_update()
# prf.read_pdb('test.pdb')
# read_param('test.param')
